#include "BTreeNode.h"
#include <iostream> // ADDED

using namespace std;

#define CS143_TRY(COMMAND)              \
do {                                    \
    RC retVal = COMMAND;                \
	if (retVal != 0) { return retVal; } \
} while (0)


//=============================================================================
//  BTLeafNode
//=============================================================================

/*
 * Read the content of the node from the page pid in the PageFile pf.
 * @param pid[IN] the PageId to read
 * @param pf[IN] PageFile to read from
 * @return 0 if successful. Return an error code if there is an error.
 */
RC BTLeafNode::read(PageId pid, const PageFile& pf)
{
	// read page into buffer
	int retVal = pf.read(pid, buffer);

	// also, populate the "node" vector for internal manipulation
	if (retVal == 0) {
		Entry tempEntry;

		node.clear();

		// update keyCount (from first 4 bytes of buffer)
		memcpy(&keyCount, buffer, sizeof(int));

		// update nextNodePtr (from second 4 bytes of buffer)
		memcpy(&nextNodePtr, buffer + sizeof(int), sizePid);

		// populate entries of vector (there are "#keyCount" entries)
		for (int i = 0; i < keyCount; ++i) {
			memcpy(&(tempEntry.rid.pid), buffer + sizeof(int) + sizePid + (i * sizeEntry), sizePid);
			memcpy(&(tempEntry.rid.sid), buffer + sizeof(int) + sizePid + sizePid + (i * sizeEntry), sizeof(int));
			memcpy(&(tempEntry.key),     buffer + sizeof(int) + sizePid + sizePid + sizeof(int) + (i * sizeEntry), sizeKey);

			node.push_back(tempEntry);
		}
		sort(node.begin(), node.end());
	}

	return retVal;
}

/*
 * Write the content of the node to the page pid in the PageFile pf.
 * @param pid[IN] the PageId to write to
 * @param pf[IN] PageFile to write to
 * @return 0 if successful. Return an error code if there is an error.
 */
RC BTLeafNode::write(PageId pid, PageFile& pf)
{
	return pf.write(pid, buffer);
}

/*
 * Return the number of keys stored in the node.
 * @return the number of keys in the node
 */
int BTLeafNode::getKeyCount()
{
	if (keyCount == -1) {
		memcpy(&keyCount, buffer, sizeof(int));
	}
	return keyCount;
}

/*
 * Insert a (key, rid) pair to the node.
 * @param key[IN] the key to insert
 * @param rid[IN] the RecordId to insert
 * @return 0 if successful. Return an error code if the node is full.
 */
RC BTLeafNode::insert(int key, const RecordId& rid)
{
	if (getKeyCount() + 1 > maxEntries) {
		return RC_NODE_FULL;
	}

	Entry e;
	e.key = key;
	e.rid = rid;

	node.push_back(e);
	sortNode();

	// update keyCount
	keyCount = node.size();

	// update our buffer
	updateBuffer();

	return 0;
}


/*
 * Insert the (key, rid) pair to the node
 * and split the node half and half with sibling.
 * The first key of the sibling node is returned in siblingKey.
 * @param key[IN] the key to insert.
 * @param rid[IN] the RecordId to insert.
 * @param sibling[IN] the sibling node to split with. This node MUST be EMPTY when this function is called.
 * @param siblingKey[OUT] the first key in the sibling node after split.
 * @return 0 if successful. Return an error code if there is an error.
 */
RC BTLeafNode::insertAndSplit(int key, const RecordId& rid,
                              BTLeafNode& sibling, int& siblingKey)
{
	if (sibling.getKeyCount() != 0)
		return RC_INVALID_ATTRIBUTE;

	// temporarily insert new entry and re-sort
	Entry e;
	e.key = key;
	e.rid = rid;

	node.push_back(e);
	sortNode();

	int curSize  = node.size();
	int sibSize  = curSize / 2; // floor(curSize/2)
	int midpoint = curSize - sibSize;

	// populate sibling's vector with 2nd half of our vector
	for (int i = midpoint; i < curSize; ++i) {
		// sibling's buffer is updated within insert
		CS143_TRY(sibling.insert(node[i].key, node[i].rid));
	}

	// return siblingKey
	siblingKey = node[midpoint].key;

	// now we can remove the pairs from our vector
	for (int i = 0; i < sibSize; ++i) {
		node.pop_back();
	}

	// update keyCount
	keyCount = node.size();

	// update our buffer
	updateBuffer();

	return 0;
}

void BTLeafNode::sortNode()
{
	::sort(node.begin(), node.end());
}

void BTLeafNode::updateBuffer()
{
	// copy vector data into buffer
	for (int i = 0; i < keyCount; ++i) {
		// copy pid into buffer
		memcpy(buffer + sizeof(int) + sizePid + (i * sizeEntry), &(node[i].rid.pid), sizePid);
		// copy sid into buffer
		memcpy(buffer + sizeof(int) + sizePid + (i * sizeEntry) + sizePid, &(node[i].rid.sid), sizeof(int));
		// copy key into buffer
		memcpy(buffer + sizeof(int) + sizePid + (i * sizeEntry) + sizePid + sizeof(int), &(node[i].key), sizeKey);
	}
	// copy keyCount into first 4 bytes of buffer
	memcpy(buffer, &keyCount, sizeof(int));

	// copy nextNodePtr into second 4 bytes of buffer
	memcpy(buffer + sizeof(int), &nextNodePtr, sizePid);
}

void BTLeafNode::clearBuffer()
{
	memset(buffer, 0, PageFile::PAGE_SIZE);
}

/*
 * Find the entry whose key value is larger than or equal to searchKey
 * and output the eid (entry number) whose key value >= searchKey.
 * Remeber that all keys inside a B+tree node should be kept sorted.
 * @param searchKey[IN] the key to search for
 * @param eid[OUT] the entry number that contains a key larger than or equality to searchKey
 * @return 0 if successful. Return an error code if there is an error.
 */
RC BTLeafNode::locate(int searchKey, int& eid)
{
	int count = 0;
	vector<Entry>::iterator it;
	it = node.begin();

	while (it != node.end()) {
		if (it->key >= searchKey) {
			eid = count;
			return 0;
		}
		it++;
		count++;
	}

	return RC_NO_SUCH_RECORD;
}

/*
 * Read the (key, rid) pair from the eid entry.
 * @param eid[IN] the entry number to read the (key, rid) pair from
 * @param key[OUT] the key from the entry
 * @param rid[OUT] the RecordId from the entry
 * @return 0 if successful. Return an error code if there is an error.
 */
RC BTLeafNode::readEntry(int eid, int& key, RecordId& rid)
{
	if (eid >= keyCount || eid < 0) {
		return RC_NO_SUCH_RECORD;
	}

	key = node[eid].key;
	rid = node[eid].rid;

	return 0;
}

/*
 * Return the pid of the next sibling node.
 * @return the PageId of the next sibling node
 */
PageId BTLeafNode::getNextNodePtr()
{
	return nextNodePtr;
}

/*
 * Set the pid of the next sibling node.
 * @param pid[IN] the PageId of the next sibling node
 * @return 0 if successful. Return an error code if there is an error.
 */
RC BTLeafNode::setNextNodePtr(PageId pid)
{
	// allow -1 as next pointer (for right-most leaf)
	if (pid < -1) {
		return RC_INVALID_PID;
	}
	else {
		nextNodePtr = pid;

		// update our buffer
		updateBuffer();

		return 0;
	}
}



//=============================================================================
//  BTNonLeafNode
//=============================================================================

/*
 * Read the content of the node from the page pid in the PageFile pf.
 * @param pid[IN] the PageId to read
 * @param pf[IN] PageFile to read from
 * @return 0 if successful. Return an error code if there is an error.
 */
RC BTNonLeafNode::read(PageId pid, const PageFile& pf)
{
	// read page into buffer
	int retVal = pf.read(pid, buffer);

	// also, populate the "node" vector for internal manipulation
	if (retVal == 0) {
		Entry tempEntry;

		node.clear();

		// update keyCount (from first 4 bytes of buffer)
		memcpy(&keyCount, buffer, sizeof(int));

		// update lhsPtr (from second 4 bytes of buffer)
		memcpy(&lhsPtr, buffer + sizeof(int), sizePid);

		// populate entries of vector (there are "#keyCount" entries)
		for (int i = 0; i < keyCount; ++i) {
			memcpy(&(tempEntry.key), buffer + sizeof(int) + sizePid + (i * sizeEntry), sizeKey);
			memcpy(&(tempEntry.pid), buffer + sizeof(int) + sizePid + sizeKey + (i * sizeEntry), sizePid);

			node.push_back(tempEntry);
		}
		sort(node.begin(), node.end());
	}

	return retVal;
}

/*
 * Write the content of the node to the page pid in the PageFile pf.
 * @param pid[IN] the PageId to write to
 * @param pf[IN] PageFile to write to
 * @return 0 if successful. Return an error code if there is an error.
 */
RC BTNonLeafNode::write(PageId pid, PageFile& pf)
{
	return pf.write(pid, buffer);
}

/*
 * Return the number of keys stored in the node.
 * @return the number of keys in the node
 */
int BTNonLeafNode::getKeyCount()
{
	if (keyCount == -1) {
		memcpy(&keyCount, buffer, sizeof(int));
	}
	return keyCount;
}

/*
 * Insert a (key, pid) pair to the node.
 * @param key[IN] the key to insert
 * @param pid[IN] the PageId to insert
 * @return 0 if successful. Return an error code if the node is full.
 */
RC BTNonLeafNode::insert(int key, PageId pid)
{
	if (getKeyCount() + 1 > maxEntries) {
		return RC_NODE_FULL;
	}

	Entry e;
	e.key = key;
	e.pid = pid;

	node.push_back(e);
	sortNode();

	// update keyCount
	keyCount = node.size();

	// update our buffer
	updateBuffer();

	return 0;
}

/*
 * Insert the (key, pid) pair to the node
 * and split the node half and half with sibling.
 * The middle key after the split is returned in midKey.
 * @param key[IN] the key to insert
 * @param pid[IN] the PageId to insert
 * @param sibling[IN] the sibling node to split with. This node MUST be empty when this function is called.
 * @param midKey[OUT] the key in the middle after the split. This key should be inserted to the parent node.
 * @return 0 if successful. Return an error code if there is an error.
 */
RC BTNonLeafNode::insertAndSplit(int key, PageId pid, BTNonLeafNode& sibling, int& midKey)
{
	if (sibling.getKeyCount() != 0)
		return RC_INVALID_ATTRIBUTE;

	// temporarily insert new entry and re-sort
	Entry e;
	e.key = key;
	e.pid = pid;

	node.push_back(e);
	sortNode();

	int curSize  = node.size();
	int midIndex = curSize / 2; // floor(size/2)

	// populate sibling's vector with 2nd half of our vector
	// sibling is guaranteed to be sorted since original vector is sorted
	for (int i = midIndex + 1; i < curSize; ++i) {
		// sibling's buffer is updated within insert
		CS143_TRY(sibling.insert(node[i].key, node[i].pid));
	}

	// return midKey
	midKey = node[midIndex].key;

	// set sibling's lhsPtr
	sibling.setLhsPtr(node[midIndex].pid);

	// now we can remove the pairs from our vector (including middle key)
	for (int i = midIndex; i < curSize; ++i) {
		node.pop_back();
	}

	// update keyCount
	keyCount = node.size();

	// update our buffer
	updateBuffer();

	return 0;
}

void BTNonLeafNode::sortNode()
{
	::sort(node.begin(), node.end());
}

void BTNonLeafNode::updateBuffer()
{
	// copy vector data into buffer
	for (int i = 0; i < keyCount; ++i) {
		// copy key into buffer
		memcpy(buffer + sizeof(int) + sizePid + (i * sizeEntry), &(node[i].key), sizeKey);
		// copy pid into buffer
		memcpy(buffer + sizeof(int) + sizePid + sizeKey + (i * sizeEntry), &(node[i].pid), sizePid);
	}

	// copy keyCount into first 4 bytes of buffer
	memcpy(buffer, &keyCount, sizeof(int));

	// copy lhsPtr into second 4 bytes of buffer
	memcpy(buffer + sizeof(int), &lhsPtr, sizePid);
}

void BTNonLeafNode::clearBuffer()
{
	memset(buffer, 0, PageFile::PAGE_SIZE);
}

/*
 * Given the searchKey, find the child-node pointer to follow and
 * output it in pid.
 * @param searchKey[IN] the searchKey that is being looked up.
 * @param pid[OUT] the pointer to the child node to follow.
 * @return 0 if successful. Return an error code if there is an error.
 */
RC BTNonLeafNode::locateChildPtr(int searchKey, PageId& pid)
{
	// find child ptr (last key that's less than searchKey)
	if (searchKey < node[0].key) {
		pid = lhsPtr;
		return 0;
	}

	int i = 0;
	while (i < node.size() && searchKey >= node[i].key) {
		i++;
	}

	pid = node[i-1].pid;

	return 0;
}

/*
 * Initialize the root node with (pid1, key, pid2).
 * @param pid1[IN] the first PageId to insert
 * @param key[IN] the key that should be inserted between the two PageIds
 * @param pid2[IN] the PageId to insert behind the key
 * @return 0 if successful. Return an error code if there is an error.
 */
RC BTNonLeafNode::initializeRoot(PageId pid1, int key, PageId pid2)
{
	Entry e;
	e.key = key;
	e.pid = pid2;

	node.push_back(e);

	lhsPtr = pid1;
	keyCount = node.size();

	// add keyCount to buffer
	memcpy(buffer, &keyCount, sizeof(int));

	// pid1 = lhsPtr. add to buffer
	memcpy(buffer + sizeof(int), &lhsPtr, sizePid);

	// add {key, pid2} to buffer
	memcpy(buffer + sizeof(int) + sizePid, &key, sizeKey);
	memcpy(buffer + sizeof(int) + sizePid + sizeKey, &pid2, sizePid);

	return 0;
}

RC BTNonLeafNode::setLhsPtr(PageId pid) {
	lhsPtr = pid;
	updateBuffer();
}
